'''
Title: Pupil: Eye Tracking Software
methods.py
Authors: Moritz Kassner & William Patera
Date: July 12, 2011
Notes: Prototype tested with 2 XBox Live Webcams on Macbook Pro OSX 10.6.7 and Ubuntu 10.10 using OpenCV 2.2
'''

import cv
import numpy as np

def extract_darkspot_ellipse_center(image,threshold):
	size = cv.GetSize(image)
	filtered = cv.CreateImage(size, 8, 1)
	gray = cv.CreateImage(size, 8, 1)
	cv.CvtColor(image, gray, cv.CV_BGR2GRAY)
	filtered = erase_specular(gray)
	
	cv.ShowImage("Eye_Gray", filtered)

	cv.InRangeS(filtered, cv.Scalar(0), cv.Scalar(threshold), filtered)
	
	cv.ShowImage("Eye_Smooth", filtered) 
	#cv.Dilate(filtered,filtered,None,1)
	cv.Erode(filtered,filtered,None,2)
	
	filtered, center = fit_ellipse(filtered)
	#this could be more efficient by using an older image
	gray = cv.CreateImage(size, 8, 1)
	cv.CvtColor(image, gray, cv.CV_BGR2GRAY)
	cv.Max(gray,filtered,filtered)
	return filtered, center
	
def erase_specular(image,threshold=180,width= 2):
	"""find specular white reflections and erode them"""
	kernel = cv.CreateStructuringElementEx(9,9,5,5, cv.CV_SHAPE_CROSS) 
	size = cv.GetSize(image)
	hi_mask = cv.CreateImage(size, 8, 1)
	hilight = cv.CreateImage(size, 8, 1)
	cv.InRangeS(image, 150, 600, hi_mask)
	cv.Erode(image,hilight,kernel,2)
	cv.Dilate(hi_mask,hi_mask,kernel,2)
	cv.Copy(hilight, image, mask=hi_mask)
	return image
  
def erase_specular_infill(image,threshold=180,width= 2):
	"""find specular white reflections and erode them"""
	#this method needs improvement
	kernel = cv.CreateStructuringElementEx(7,7,4,4, cv.CV_SHAPE_RECT) 
	size = cv.GetSize(image)
	hi_mask = cv.CreateImage(size, 8, 1)
	hilight = cv.CreateImage(size, 8, 1)
	cv.Erode(image,image,None,1)

	cv.InRangeS(image, 150, 600, hi_mask)
	cv.Inpaint(image, hi_mask, hilight, 2, cv.CV_INPAINT_TELEA) 
	cv.Erode(hilight,hilight,kernel,1)
	cv.Dilate(hi_mask,hi_mask,kernel,1)
	cv.Copy(hilight, image, mask=hi_mask)

	return image
 
def contour_iterator(contour):
		while contour:
			yield contour
			contour = contour.h_next()
			
def fit_ellipse(image):
	"""Fit ellipse around pupil"""
	stor = cv.CreateMemStorage()
	cont = cv.FindContours(image,stor,cv.CV_RETR_LIST,cv.CV_CHAIN_APPROX_NONE,(0, 0))
	max_center = [None,0]
	for c in contour_iterator(cont):
	# Number of points must be more than or equal to 6 for cv.FitEllipse2
		if len(c) >= 80:
			# Copy the contour into an array of (x,y)s
			PointArray2D32f = cv.CreateMat(1, len(c), cv.CV_32FC2)
			for (i, (x, y)) in enumerate(c):
				PointArray2D32f[0, i] = (x, y)
			# Fits ellipse to current contour.
			(center, size, angle) = cv.FitEllipse2(PointArray2D32f)
		   
			if size > max_center[1]:
				max_center[0] = center
				max_center[1] = size
			if True:
				 # Draw the current contour in gray
				gray = cv.CV_RGB(255, 255, 255)
				cv.DrawContours(image, c, gray, gray,0,1,8,(0,0))
			   
				# Convert ellipse data from float to integer representation.
				center = (cv.Round(center[0]), cv.Round(center[1]))
				size = (cv.Round(size[0] * 0.5), cv.Round(size[1] * 0.5))
				angle = angle
   
				color = cv.CV_RGB(255,255,255)
				cv.Ellipse(image, center, size,
						  angle, 0, 360,
						  color, 1, cv.CV_AA, 0)

	if max_center[0]:
		major = max(max_center[1])
		minor = min(max_center[1])
		ratio = major/minor 
		#print ratio
		if ratio >3:
			print "blink"
			max_center[0] = None
	return image, max_center[0]
	

def chessboard(img):
	"""this finds the chessboard corner of an image and draws them onto it
	return the image and the center of the chessboard as (x,y)
	if no chessboad found it return the img and None"""
	
	row = 3
	col = 3 
	
	#corners --> cv2 is for opencv2.3
	#corners = cv2.findChessboardCorners(cv.GetMat(img), (col,row), cv.CV_CALIB_CB_ADAPTIVE_THRESH + cv.CV_CALIB_CB_NORMALIZE_IMAGE + cv2.CALIB_CB_FAST_CHECK)
	corners = cv.FindChessboardCorners(cv.GetMat(img), (col,row), cv.CALIB_CB_ADAPTIVE_THRESH + cv.CALIB_CB_NORMALIZE_IMAGE + cv.CALIB_CB_FAST_CHECK);
	if corners[0]:
		corners = np.array(corners[1]) #array of all corners
		# copy the corners to corner list - this will not be necessary in future iterations, but is necessary for current debugging
		# make a new list 'corner_lst' with x and y types
		# sort corner_lst by x and y values
		# this method will be used to cycle through all checkerboard intersections
		# -- the eye will look at these points in sequence
		# ideally this will yeild a more robust/standard calibration routine
		
		corner_lst = corners[:]
		corner_lst.dtype = np.dtype({'names': ['x','y',],'formats': ['f8', 'f8']})
		corner_lst.sort(axis=0, order = "y")
		np.reshape(corner_lst, (row,col))
		corner_lst.sort(axis=1, order = "x")
		corner_lst.shape= (row*col)
		print 'first corner', corner_lst[0]
		pt = (int(corner_lst[0][0]), int(corner_lst[0][1]))
		cv.Circle(img, pt, 3, cv.CV_RGB(0, 150, 255), thickness=1, lineType=cv.CV_AA, shift=0)
	
		#old method, draws the center of the checkerboard and returns it for calibration process
		mean = corners.sum(0)/corners.shape[0]
		mean = (int(mean[0]),int(mean[1]))
		cv.Circle(img, mean, 3, cv.CV_RGB(255, 0, 0), thickness=1, lineType=cv.CV_AA, shift=0)
		return img,mean
	else:
		return img,None

############################## - old methods - ##############################	
def chessboard_old(img):
	"""this finds the chessboard corner of an image and draws them onto it
	return the image and the center of the chessboard as (x,y)
	if no chessboad found it return the img and None"""
	#corners --> cv2 is for opencv2.3
	#corners = cv2.findChessboardCorners(cv.GetMat(img), (3,3), cv.CV_CALIB_CB_ADAPTIVE_THRESH + cv.CV_CALIB_CB_NORMALIZE_IMAGE + cv2.CALIB_CB_FAST_CHECK)
	corners = cv.FindChessboardCorners(cv.GetMat(img), (3,3), cv.CALIB_CB_ADAPTIVE_THRESH + cv.CALIB_CB_NORMALIZE_IMAGE + cv.CALIB_CB_FAST_CHECK);
	if corners[0]:
		for corner in corners[1]:
			corner_x = int(corner[0])
			corner_y = int(corner[1])
			corner = (corner_x, corner_y)
			cv.Circle(img, corner, 3, cv.CV_RGB(255, 255, 255), thickness=1, lineType=cv.CV_AA, shift=0)
		corners = np.array(corners[1])
		mean = corners.sum(0)/corners.shape[0]
		mean = (int(mean[0]),int(mean[1]))
		cv.Circle(img,mean, 3, cv.CV_RGB(255, 0, 0), thickness=1, lineType=cv.CV_AA, shift=0)
		return img,mean
	else:
		return img,None
	
def extract_darkspot_moment_old(image,threshold):
	size = cv.GetSize(image)
	filtered = cv.CreateImage(size, 8, 1)
	hi_mask = cv.CreateImage(size, 8, 1)
	hilight = cv.CreateImage(size, 8, 1)
	cv.CvtColor(image, filtered, cv.CV_BGR2GRAY)
	cv.Erode(filtered,hilight,None,1)
	cv.InRangeS(filtered, 180, 300, hi_mask)
	cv.Erode(filtered,hilight,None,2)
	cv.Dilate(hi_mask,hi_mask,None,2)
	cv.Copy(hilight, filtered, mask=hi_mask)
	cv.ShowImage("Eye_Gray", filtered)
	cv.InRangeS(filtered, cv.Scalar(0), cv.Scalar(threshold), filtered)
	cv.Smooth(filtered, filtered, smoothtype=cv.CV_BLUR_NO_SCALE, param1=2, param2=0, param3=0, param4=0)
	cv.ShowImage("Eye_Smooth", filtered) 
	cv.Dilate(filtered,filtered,None,1)
	cv.Erode(filtered,filtered,None,2)
	return filtered, cv.Moments(cv.GetMat(filtered))
	
def extract_colorspot_moment(image,threshold):
	size = cv.GetSize(image)
	filtered = cv.CreateImage(size, 8, 1)
	
	B = cv.CreateImage(size, 8, 1)
	G = cv.CreateImage(size, 8, 1)
	R = cv.CreateImage(size, 8, 1)
	HSV = cv.CreateImage(size, 8, 3)
	H = cv.CreateImage(size, 8, 1)
	S = cv.CreateImage(size, 8, 1)
	V = cv.CreateImage(size, 8, 1)
	cv.CvtColor(image, HSV, cv.CV_BGR2HSV)
	cv.Split(HSV, H,S,V, None)
	cv.Split(image,B,G,R,None)
	cv.InRangeS(B,120,255,B)
	cv.InRangeS(S,120,255,S)  
	cv.And(S,B,filtered)  
	cv.Erode(filtered,filtered,None,1)
	return filtered, cv.Moments(cv.GetMat(filtered))
		  
def get_pos_area(moments):
	moment_10 = cv.GetSpatialMoment(moments, 1, 0)
	moment_01 = cv.GetSpatialMoment(moments, 0, 1) 
	area = cv.GetCentralMoment(moments, 0, 0)
	if area:
		x = moment_10/area
		y = moment_01/area
		return x,y,area
	else:
		return None,None,None

#useless but required!
def UpdateVar(val):
	global p_threshold
	p_threshold = max(val,30)